# -*- coding: UTF-8 -*-
from config.getConfig import *
import gol
import time
import os
import json


def predict(model):
    # PREDICT = False
    while True:
        while gol.get_queue_size('pic_line') > 0 and not gol.get_value('predict_lock'):
            ####
            start = time.clock()
            ####

            # PREDICT = True
            gol.set_value('predict_lock', True)
            data = gol.get_queue_value('pic_line')
            result = predict_threshold(model, data['Path'], getConfigFloat("predict.ini",
                                                                            "predict",
                                                                            "threshold"))
            result_with_distance = compute_distance(result)
            # print(json.dumps(result_with_distance))
            # print(type(json.dumps(result_with_distance)))
            num_people = len(result_with_distance)
            num_people = int(data['MaxNum']) if num_people > int(data['MaxNum']) else num_people
            predict_data = {'BuildName': data['BuildName'],
                            'ClassName': data['ClassName'],
                            'Time': data['Time'],
                            'Num': num_people,
                            'MaxNum': data['MaxNum'],
                            'Bbox': json.dumps({"data": result_with_distance})}
            gol.put_queue_value('predict_line', predict_data)
            print("\n-- Time:{} --".format(data['Time']))
            print('预测完成! 预测人数：{} 教室最大人数：{}'.format(num_people, data['MaxNum']))
            print('待预测图片数量：{}'.format(gol.get_queue_size('pic_line')))
            # gol.put_queue_value('inf_line', '预测完成! 预测人数：{} 教室最大人数：{}'.format(num_people, data['MaxNum']))
            os.remove(data['Path'])
            # PREDICT = False
            gol.set_value('predict_lock', False)

            ###
            lapse = time.clock() - start
            gol.set_value('predict_time', gol.get_value('predict_time') + lapse)
            gol.set_value('predict_times', gol.get_value('predict_times') + 1)
            print("预测平均时间{}".format(gol.get_value('predict_time') / gol.get_value('predict_times')))
            ###
        time.sleep(0.01)


def predict_threshold(model, pic_path, threshold):
    """
    预测图片，过滤置信度低于阈值的目标
    :param model: paddle模型
    :param pic_path: 待预测的图片的地址
    :param threshold: 阈值
    :return: 返回一个元素为字典的列表，字典包含category_id、bbox、score、category
    """
    result = model.predict(pic_path)
    result_threshold = []
    for bbox in result:
        if float(bbox['score']) >= threshold:
            result_threshold.append(bbox)
    return result_threshold


def compute_distance(result):
    """
    计算预测出的目标之间的距离
    :param result: predict_threshold函数返回的result
    :return: 与result结构相同的list，字典中添加了一个新的键值对，is_too_close:True/False True为离得太近
    """
    mini_distance = getConfigFloat("predict.ini", "distance", "distance")
    for i in range(len(result)):
        if 'is_too_close' in result[i]:
            continue
        flag = False
        for j in range(i + 1, len(result)):
            point_1 = [result[i]['bbox'][0] + result[i]['bbox'][2] / 2, result[i]['bbox'][1] + result[i]['bbox'][3] / 2]
            point_2 = [result[j]['bbox'][0] + result[j]['bbox'][2] / 2, result[j]['bbox'][1] + result[j]['bbox'][3] / 2]
            distance = (pow((point_1[0] - point_2[0]), 2) + pow((point_1[1] - point_2[1]), 2)) ** 0.5
            if distance < mini_distance:
                flag = True
                result[i]['is_too_close'] = flag
                result[j]['is_too_close'] = flag
                break
        if not flag:
            result[i]['is_too_close'] = flag
        del result[i]['category_id']
        del result[i]['category']
    return result




